import struct
import pcap
import socket
import binascii
import string

# decode IP Packet
def decode_ip_packet(s):
    d={}
    d['version']=(ord(s[0]) & 0xf0) >> 4
    d['header_len']=ord(s[0]) & 0x0f
    d['tos']=ord(s[1])
    d['total_len']=socket.ntohs(struct.unpack('H',s[2:4])[0])
    d['id']=socket.ntohs(struct.unpack('H',s[4:6])[0])
    d['flags']=(ord(s[6]) & 0xe0) >> 5
    d['fragment_offset']=socket.ntohs(struct.unpack('H',s[6:8])[0] & 0x1f)
    d['ttl']=ord(s[8])
    d['protocol']=ord(s[9])
    d['checksum']=socket.ntohs(struct.unpack('H',s[10:12])[0])
    d['source_address']=pcap.ntoa(struct.unpack('i',s[12:16])[0])
    d['destination_address']=pcap.ntoa(struct.unpack('i',s[16:20])[0])
    if d['header_len']>5:
        d['options']=s[20:4*(d['header_len']-5)]
    else:
        d['options']=None
    d['data']=s[4*d['header_len']:]
    return d


# decode TCP packet
def decode_tcp_packet(s):
    d = {}
    d['src_port']=socket.ntohs(struct.unpack('H',s[0:2])[0])
    d['dst_port']=socket.ntohs(struct.unpack('H',s[2:4])[0])
    d['seq_num']=socket.ntohl(struct.unpack('i',s[4:8])[0])
    d['ack_num']=socket.ntohl(struct.unpack('i',s[8:12])[0])
    d['header_len']=(ord(s[12]) & 0xf0) >> 4
    d['flags'] = ord(s[13]) & 0x3f
    d['window'] = socket.ntohs(struct.unpack('H',s[14:16])[0])
    d['checksum'] = socket.ntohs(struct.unpack('H',s[16:18])[0])
    #d['urgent'] 
    d['data']=s[4*d['header_len']:]

    return d

# decode UDP packet
def decode_udp_packet(s):
    d = {}
    d['src_port']=socket.ntohs(struct.unpack('H',s[0:2])[0])
    d['dst_port']=socket.ntohs(struct.unpack('H',s[2:4])[0])
    d['length']=socket.ntohs(struct.unpack('H',s[4:6])[0])
    d['checksum']=socket.ntohs(struct.unpack('H',s[6:8])[0])
    d['data']=s[8:]
    return d

# decode ICMP_packet
def decode_icmp_packet(s):
    d = {}
    d['type'] = ord(s[0])
    d['code'] = ord(s[1])
    return d

def decode_pseudo_packet(s):
    """ decode IP header (src ip, dst ip, proto)
    tcp/udp src port, dst port
    payload : 10 bytes:(if exist?)
    """
    d = {}
    # IP header
    hdr_len = ord(s[0]) & 0x0f
    proto = ord(s[9])
    d['proto'] = proto
    d['src_ip'] = pcap.ntoa(struct.unpack('i',s[12:16])[0])
    d['dst_ip'] = pcap.ntoa(struct.unpack('i',s[16:20])[0])
    data = s[4*hdr_len:]
    
    # TCP/UDP header
    d['src_port'] = socket.ntohs(struct.unpack('H',data[0:2])[0])
    d['dst_port'] = socket.ntohs(struct.unpack('H',data[2:4])[0])
    if proto == 6:
        # TCP
        tcp_hdr_len = (ord(data[12]) & 0xf0) >> 4
        d['payload'] = data[4*tcp_hdr_len:] 
    elif proto == 17:
        # UDP
        d['payload'] = data[8:]
    elif proto == 1:
        # ICMP
        d['src_port'] = -1
        d['dst_port'] = -1
        d['payload'] = data[hdr_len:]
    else:
        d['src_port'] = -1
        d['dst_port'] = -1
        d['payload'] = data[hdr_len:]
    return d

def format_dns(s):
    """ parse DNS payload
    s is payload (layer 7 data)
    """
    d = {}
    if len(s) <= 12:
        print "payload:%s" % s
        print "Not DNS Header"
        d['query']="error"
        d['querytype']=-1
        print d
        return d
    data = s[12:]
    count = int(ord(data[0]))
    query = ""
    index = 0
    while count != 0:
        next_index = index + count + 1
        query = query + data[index+1:next_index]
        index = next_index
        count = int(ord(data[index]))
        if count == 0:
            break
        query = query + "." 
    d['query'] = query
    d['querytype'] = socket.ntohs(struct.unpack('H',data[index+1:index+3])[0])
    return d

# dump payload
def dumphex(s):
  bytes = map(lambda x: '%.2x' % x, map(ord, s))
  output = ''
  i = 0
  for i in xrange(0,len(bytes)/16):
    output = string.join(bytes[i*16:(i+1)*16],' ')
  output = output +  string.join(bytes[(i+1)*16:],' ')
  return output + "\n"

def dumpstring(s):
  bytes = map(lambda x: '%.2x' % x, map(ord, s))
  output = ''
  for i in xrange(0, len(bytes)):
    output = output + binascii.a2b_hex(bytes[i])

  return output + "\n"

